//
// Created by inter on 19-10-17.
//

#include <iostream>
#include <cstring>

#include "gpu_info.h"
#include "nvml.h"

bool GPUInfo::init(){
    return nvmlInit() == NVML_SUCCESS;
}

bool GPUInfo::shunDown(){
    return nvmlShutdown() == NVML_SUCCESS;
}

bool GPUInfo::getInformation(GPUI_Info& info) {
    int ret = (int) nvmlDeviceGetCount(&info.gpu_count);
    if (ret != NVML_SUCCESS) {
        std::cout << "ERROR: nvmlDeviceGetCount = "<< ret << std::endl;
        return false;
    }

    char buf[128] = {0};
    nvmlDevice_t device_handle;
    unsigned int index = 0;
    while(index < info.gpu_count) {
        ret = nvmlDeviceGetHandleByIndex(index, &device_handle);
        if (ret != NVML_SUCCESS) {
            std::cout << "ERROR: nvmlDeviceGetHandleByIndex = "<< ret << std::endl;
            return false;
        }

        if (index == 0) {
            ret = nvmlDeviceGetBoardPartNumber(device_handle, buf, 128);
            if (ret != NVML_SUCCESS) {
                std::cout << "ERROR: getGPUBoardPartNumber = "<< ret << std::endl;
                return false;
            }
            info.board_serial = buf;
        }

        vector<string> item;
        memset(buf, 0, 128);
        ret = nvmlDeviceGetUUID(device_handle, buf, 128);
        if (ret == NVML_SUCCESS){
            item.emplace_back(buf);
            // 读取gpu的name
            memset(buf, 0, 128);
            ret = nvmlDeviceGetName(device_handle, buf, 128);
            if (ret == NVML_SUCCESS) {
                item.emplace_back(string(buf));
            } else{
                std::cout << "ERROR: nvmlDeviceGetName = "<< ret << std::endl;
            }
        } else {
            std::cout << "ERROR: nvmlDeviceGetUUID = "<< ret << std::endl;
        }

        info.gpu_list.push_back(item);
        index++;
    }

    return true;
}

bool GPUInfo::getGPUCount(unsigned int& count){
    int ret = (int) nvmlDeviceGetCount(&count);
    if (ret != NVML_SUCCESS) {
        std::cout << "ERROR: nvmlDeviceGetCount = "<< ret << std::endl;
        return false;
    }
    return true;
}

bool GPUInfo::getGPUBName(unsigned int index, string& name, string& uuid) {
    nvmlDevice_t device_handle;
    int ret = nvmlDeviceGetHandleByIndex(index, &device_handle);
    if (ret != NVML_SUCCESS) {
        std::cout << "ERROR: nvmlDeviceGetHandleByIndex = "<< ret << std::endl;
        return false;
    }

    char buf[128] = {0};
    ret = nvmlDeviceGetUUID(device_handle, buf, 128);
    if (ret == NVML_SUCCESS){
        uuid = string(buf);
        // 读取gpu的name
        memset(buf, 0, 128);

        ret = nvmlDeviceGetName(device_handle, buf, 32);
        if (ret == NVML_SUCCESS) {
            name = string(buf);
        } else{
            std::cout << "ERROR: nvmlDeviceGetName = "<< ret << std::endl;
        }
        return true;
    } else {
        std::cout << "ERROR: nvmlDeviceGetUUID = "<< ret << std::endl;
        return false;
    }
}

